import time
from typing import Optional

import jwt
from pydantic import BaseModel
from starlette.requests import Request
from starlette.responses import JSONResponse
from starlette.types import ASGIApp, Receive, Scope, Send

from fastapi_template.config import app_config
from fastapi_template.webapp.responses import CommonResp


class TokenData(BaseModel):
    user_id: str
    exp: int


class AuthFailed(Exception):
    pass


class UserBlocked(Exception):
    pass


def check_token(token):
    if not token:
        return False
    jwt_secret = app_config.JWT_SECRET
    jwt_algorithm = app_config.JWT_ALGORITHM

    try:
        token_data_dict = jwt.decode(token,
                                     jwt_secret,
                                     algorithms=[jwt_algorithm])
        token_data = TokenData.model_validate(token_data_dict)
    except jwt.exceptions.InvalidTokenError:
        return False

    return token_data


def make_token(user_id: str) -> str:
    exp = int(time.time()) + app_config.JWT_MAX_AGE
    token_data = TokenData(user_id=user_id, exp=exp)
    token = jwt.encode(
        token_data.model_dump(),
        app_config.JWT_SECRET,
        algorithm=app_config.JWT_ALGORITHM,
    )
    return token


def fetch_csrf_token(request: Request) -> str | None:
    for k, v in request.headers.items():
        if isinstance(k, str) and (k.lower() == "x-csrf-token"):
            return v
    return None


def fetch_jwt_token(request: Request) -> str | None:
    jwt_name = app_config.JWT_NAME
    jwt_token = request.cookies.get(jwt_name, "")
    if not jwt_token:
        jwt_token = request.headers.get("Authorization", "")
        if jwt_token.startswith("Bearer "):
            jwt_token = jwt_token[7:]
    if not jwt_token:
        jwt_token = request.query_params.get("token", "")
    return jwt_token


class AuthCheckResult(BaseModel):
    token_data: Optional[TokenData]

    def get_user_id(self) -> str | None:
        if self.token_data is None:
            return None
        return self.token_data.user_id

    def must_get_user_id(self) -> str:
        user_id = self.get_user_id()
        if user_id is None:
            raise AuthFailed("用户未登录")
        return user_id


class AuthChecker:

    def __init__(self, check_csrf: bool = True, check_login: bool = True):
        self.check_csrf = check_csrf
        self.check_login = check_login

    async def __call__(self, request: Request) -> AuthCheckResult:
        token_data = check_token(fetch_jwt_token(request))
        if not token_data:
            raise AuthFailed("用户未登录")
        return AuthCheckResult(token_data=token_data)


class HandleAuthFailedMiddleware:

    def __init__(self, app: ASGIApp) -> None:
        self.app = app

    async def __call__(self, scope: Scope, receive: Receive,
                       send: Send) -> None:
        try:
            await self.app(scope, receive, send)
        except AuthFailed:
            resp = JSONResponse(
                status_code=401,
                content=CommonResp.fail(
                    message="Unauthorized",
                    code=401).model_dump(exclude_none=True),
            )
            jwt_name = app_config.JWT_NAME
            resp.delete_cookie(jwt_name)
            await resp(scope, receive, send)
        except UserBlocked:
            resp = JSONResponse(
                status_code=499,
                content=CommonResp.fail(
                    message="暂时无法使用服务,请您联系客服",
                    code=499).model_dump(exclude_none=True),
            )
            await resp(scope, receive, send)
